#include "scepter_lerobot_driver/lerobot_hardware_interface.hpp"

#include <chrono>
#include <cmath>
#include <limits>
#include <memory>
#include <vector>

#include "hardware_interface/types/hardware_interface_type_values.hpp"
#include "rclcpp/rclcpp.hpp"

namespace scepter_lerobot_driver
{

hardware_interface::CallbackReturn LerobotHardwareInterface::on_init(
  const hardware_interface::HardwareInfo & info)
{
  if (
    hardware_interface::SystemInterface::on_init(info) !=
    hardware_interface::CallbackReturn::SUCCESS)
  {
    return hardware_interface::CallbackReturn::ERROR;
  }

  // Read parameters from URDF
  robot_ip_ = info_.hardware_parameters.count("robot_ip") ? 
    info_.hardware_parameters.at("robot_ip") : "localhost";
  config_file_ = info_.hardware_parameters.count("config_file") ? 
    info_.hardware_parameters.at("config_file") : "";
  communication_timeout_ = std::stod(
    info_.hardware_parameters.count("communication_timeout") ? 
    info_.hardware_parameters.at("communication_timeout") : "5.0");

  // Initialize joint info
  joint_names_.clear();
  hw_commands_.clear();
  hw_positions_.clear();
  hw_velocities_.clear();

  for (const hardware_interface::ComponentInfo & joint : info_.joints)
  {
    // Check joint has position command interface
    if (joint.command_interfaces.size() != 1)
    {
      RCLCPP_FATAL(
        rclcpp::get_logger("LerobotHardwareInterface"),
        "Joint '%s' has %zu command interfaces found. 1 expected.", 
        joint.name.c_str(), joint.command_interfaces.size());
      return hardware_interface::CallbackReturn::ERROR;
    }

    if (joint.command_interfaces[0].name != hardware_interface::HW_IF_POSITION)
    {
      RCLCPP_FATAL(
        rclcpp::get_logger("LerobotHardwareInterface"),
        "Joint '%s' have '%s' command interface. '%s' expected.", 
        joint.name.c_str(), joint.command_interfaces[0].name.c_str(),
        hardware_interface::HW_IF_POSITION);
      return hardware_interface::CallbackReturn::ERROR;
    }

    // Check joint has position and velocity state interfaces
    if (joint.state_interfaces.size() != 2)
    {
      RCLCPP_FATAL(
        rclcpp::get_logger("LerobotHardwareInterface"),
        "Joint '%s' has %zu state interface. 2 expected.", 
        joint.name.c_str(), joint.state_interfaces.size());
      return hardware_interface::CallbackReturn::ERROR;
    }

    joint_names_.push_back(joint.name);
    hw_commands_.push_back(std::numeric_limits<double>::quiet_NaN());
    hw_positions_.push_back(std::numeric_limits<double>::quiet_NaN());
    hw_velocities_.push_back(std::numeric_limits<double>::quiet_NaN());
  }

  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Initialized hardware interface for %zu joints", joint_names_.size());

  hardware_connected_ = false;

  return hardware_interface::CallbackReturn::SUCCESS;
}

hardware_interface::CallbackReturn LerobotHardwareInterface::on_configure(
  const rclcpp_lifecycle::State & /*previous_state*/)
{
  // Create node for communication
  node_ = rclcpp::Node::make_shared("lerobot_hardware_interface_node");

  // Create subscriber for joint states from lerobot bridge
  joint_state_subscriber_ = node_->create_subscription<sensor_msgs::msg::JointState>(
    "/lerobot/joint_states", 10,
    std::bind(&LerobotHardwareInterface::joint_state_callback, this, std::placeholders::_1));

  // Create publisher for joint commands to lerobot bridge
  command_publisher_ = node_->create_publisher<sensor_msgs::msg::JointState>(
    "/lerobot/joint_commands", 10);

  // Create service client for enabling/disabling robot
  enable_service_client_ = node_->create_client<std_srvs::srv::SetBool>("/lerobot/enable_robot");

  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Hardware interface configured successfully");

  return hardware_interface::CallbackReturn::SUCCESS;
}

std::vector<hardware_interface::StateInterface>
LerobotHardwareInterface::export_state_interfaces()
{
  std::vector<hardware_interface::StateInterface> state_interfaces;

  for (std::size_t i = 0; i < joint_names_.size(); i++)
  {
    state_interfaces.emplace_back(hardware_interface::StateInterface(
      joint_names_[i], hardware_interface::HW_IF_POSITION, &hw_positions_[i]));
    state_interfaces.emplace_back(hardware_interface::StateInterface(
      joint_names_[i], hardware_interface::HW_IF_VELOCITY, &hw_velocities_[i]));
  }

  return state_interfaces;
}

std::vector<hardware_interface::CommandInterface>
LerobotHardwareInterface::export_command_interfaces()
{
  std::vector<hardware_interface::CommandInterface> command_interfaces;

  for (std::size_t i = 0; i < joint_names_.size(); i++)
  {
    command_interfaces.emplace_back(hardware_interface::CommandInterface(
      joint_names_[i], hardware_interface::HW_IF_POSITION, &hw_commands_[i]));
  }

  return command_interfaces;
}

hardware_interface::CallbackReturn LerobotHardwareInterface::on_activate(
  const rclcpp_lifecycle::State & /*previous_state*/)
{
  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Activating hardware interface...");

  // Connect to robot
  if (!connect_to_robot())
  {
    RCLCPP_ERROR(
      rclcpp::get_logger("LerobotHardwareInterface"),
      "Failed to connect to robot");
    return hardware_interface::CallbackReturn::ERROR;
  }

  // Initialize commands to current positions
  for (std::size_t i = 0; i < hw_commands_.size(); i++)
  {
    if (std::isnan(hw_positions_[i]))
    {
      hw_commands_[i] = 0.0;
      hw_positions_[i] = 0.0;
    }
    else
    {
      hw_commands_[i] = hw_positions_[i];
    }
  }

  hardware_connected_ = true;
  last_read_time_ = std::chrono::steady_clock::now();

  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Hardware interface activated successfully");

  return hardware_interface::CallbackReturn::SUCCESS;
}

hardware_interface::CallbackReturn LerobotHardwareInterface::on_deactivate(
  const rclcpp_lifecycle::State & /*previous_state*/)
{
  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Deactivating hardware interface...");

  // Disable robot motors
  if (enable_service_client_->wait_for_service(std::chrono::seconds(1)))
  {
    auto request = std::make_shared<std_srvs::srv::SetBool::Request>();
    request->data = false;
    enable_service_client_->async_send_request(request);
  }

  hardware_connected_ = false;

  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Hardware interface deactivated");

  return hardware_interface::CallbackReturn::SUCCESS;
}

hardware_interface::return_type LerobotHardwareInterface::read(
  const rclcpp::Time & /*time*/, const rclcpp::Duration & /*period*/)
{
  if (!hardware_connected_)
  {
    return hardware_interface::return_type::ERROR;
  }

  // Spin node to process callbacks
  rclcpp::spin_some(node_);

  // Check for communication timeout
  auto now = std::chrono::steady_clock::now();
  auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now - last_read_time_);
  
  if (duration.count() > communication_timeout_ * 1000)
  {
    RCLCPP_WARN_THROTTLE(
      rclcpp::get_logger("LerobotHardwareInterface"),
      *node_->get_clock(), 1000,
      "Communication timeout: no data received for %ld ms", duration.count());
  }

  return hardware_interface::return_type::OK;
}

hardware_interface::return_type LerobotHardwareInterface::write(
  const rclcpp::Time & /*time*/, const rclcpp::Duration & /*period*/)
{
  if (!hardware_connected_)
  {
    return hardware_interface::return_type::ERROR;
  }

  // Send commands to robot
  if (!send_commands_to_robot())
  {
    RCLCPP_ERROR_THROTTLE(
      rclcpp::get_logger("LerobotHardwareInterface"),
      *node_->get_clock(), 1000,
      "Failed to send commands to robot");
    return hardware_interface::return_type::ERROR;
  }

  return hardware_interface::return_type::OK;
}

bool LerobotHardwareInterface::connect_to_robot()
{
  // Wait for lerobot bridge service
  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Waiting for lerobot bridge services...");

  if (!enable_service_client_->wait_for_service(std::chrono::seconds(10)))
  {
    RCLCPP_ERROR(
      rclcpp::get_logger("LerobotHardwareInterface"),
      "Lerobot bridge service not available");
    return false;
  }

  // Enable robot motors
  auto request = std::make_shared<std_srvs::srv::SetBool::Request>();
  request->data = true;
  
  auto future = enable_service_client_->async_send_request(request);
  
  if (rclcpp::spin_until_future_complete(node_, future, std::chrono::seconds(5)) != 
      rclcpp::FutureReturnCode::SUCCESS)
  {
    RCLCPP_ERROR(
      rclcpp::get_logger("LerobotHardwareInterface"),
      "Failed to enable robot");
    return false;
  }

  auto response = future.get();
  if (!response->success)
  {
    RCLCPP_ERROR(
      rclcpp::get_logger("LerobotHardwareInterface"),
      "Robot enable service returned failure");
    return false;
  }

  RCLCPP_INFO(
    rclcpp::get_logger("LerobotHardwareInterface"),
    "Successfully connected to robot");

  return true;
}

void LerobotHardwareInterface::joint_state_callback(
  const sensor_msgs::msg::JointState::SharedPtr msg)
{
  // Update hardware positions and velocities from robot feedback
  for (std::size_t i = 0; i < joint_names_.size() && i < msg->position.size(); i++)
  {
    hw_positions_[i] = msg->position[i];
    
    if (i < msg->velocity.size())
    {
      hw_velocities_[i] = msg->velocity[i];
    }
    else
    {
      hw_velocities_[i] = 0.0;
    }
  }

  last_read_time_ = std::chrono::steady_clock::now();
}

bool LerobotHardwareInterface::send_commands_to_robot()
{
  // Create and publish joint command message
  auto msg = sensor_msgs::msg::JointState();
  msg.header.stamp = node_->now();
  msg.name = joint_names_;
  msg.position = hw_commands_;

  command_publisher_->publish(msg);

  return true;
}

}  // namespace scepter_lerobot_driver

#include "pluginlib/class_list_macros.hpp"

PLUGINLIB_EXPORT_CLASS(
  scepter_lerobot_driver::LerobotHardwareInterface, 
  hardware_interface::SystemInterface)